//! Generic monomorphization pass.
//!
//! This module walks the AST, finds explicit generic invocations (currently
//! limited to free function calls with concrete type arguments), asks the
//! monomorphizer to materialize concrete versions, and rewrites the call sites
//! to reference the monomorphic symbols. Newly generated declarations are
//! appended to the program so downstream phases treat them like regular items.

use crate::ast::*;
use crate::error::{CompilerError, Location, Result};
use crate::type_check::monomorphizer::Monomorphizer;

/// Execute the monomorphization pass over the given program.
pub fn monomorphize_program(program: &mut Program) -> Result<()> {
    let mut visitor = MonomorphizationVisitor::new();
    visitor.monomorphizer.collect_templates(program);
    visitor.process_program(program)?;

    let pending_structs = visitor.monomorphizer.take_pending_struct_instantiations();
    for struct_decl in pending_structs {
        program.declarations.push(Declaration::Struct(struct_decl));
    }

    let pending_functions = visitor.monomorphizer.take_pending_function_instantiations();
    for func in pending_functions {
        program.declarations.push(Declaration::Function(func));
    }

    Ok(())
}

struct MonomorphizationVisitor {
    monomorphizer: Monomorphizer,
}

impl MonomorphizationVisitor {
    fn new() -> Self {
        Self {
            monomorphizer: Monomorphizer::new(),
        }
    }

    fn process_program(&mut self, program: &mut Program) -> Result<()> {
        for declaration in &mut program.declarations {
            self.process_declaration(declaration)?;
        }
        Ok(())
    }

    fn process_declaration(&mut self, declaration: &mut Declaration) -> Result<()> {
        match declaration {
            Declaration::Function(func) => self.process_function(func),
            Declaration::Struct(struct_decl) => {
                for method in &mut struct_decl.methods {
                    self.process_function(method)?;
                }
                Ok(())
            }
            Declaration::Impl(impl_decl) => {
                for method in &mut impl_decl.methods {
                    self.process_function(method)?;
                }
                Ok(())
            }
            Declaration::TraitImpl(trait_impl) => {
                for method in &mut trait_impl.methods {
                    self.process_function(method)?;
                }
                Ok(())
            }
            Declaration::Enum(_)
            | Declaration::Trait(_)
            | Declaration::Variable(_)
            | Declaration::Constant(_)
            | Declaration::Type(_)
            | Declaration::Import(_) => Ok(()),
        }
    }

    fn process_function(&mut self, func: &mut FunctionDecl) -> Result<()> {
        if !func.generics.is_empty() {
            return Ok(()); // Skip generic templates; they will be instantiated later.
        }

        for param in &mut func.parameters {
            self.process_type(&mut param.param_type)?;
        }
        if let Some(ret_type) = &mut func.return_type {
            self.process_type(ret_type)?;
        }

        if let Some(body) = &mut func.body {
            self.process_block(body)?;
        }
        Ok(())
    }

    fn process_block(&mut self, block: &mut Block) -> Result<()> {
        for statement in &mut block.statements {
            self.process_statement(statement)?;
        }
        Ok(())
    }

    fn process_statement(&mut self, statement: &mut Statement) -> Result<()> {
        match statement {
            Statement::Expression(expr) => self.process_expression(expr),
            Statement::VariableDecl(decl) => {
                if let Some(init) = &mut decl.initializer {
                    self.process_expression(init)?;
                }
                if let Some(var_type) = &mut decl.var_type {
                    self.process_type(var_type)?;
                }
                Ok(())
            }
            Statement::ConstantDecl(decl) => {
                self.process_expression(&mut decl.value)?;
                if let Some(const_type) = &mut decl.const_type {
                    self.process_type(const_type)?;
                }
                Ok(())
            }
            Statement::Assignment(assign) => {
                self.process_assignment_target(&mut assign.target)?;
                self.process_expression(&mut assign.value)
            }
            Statement::Return(ret) => {
                if let Some(value) = &mut ret.value {
                    self.process_expression(value)?;
                }
                Ok(())
            }
            Statement::If(if_stmt) => {
                self.process_expression(&mut if_stmt.condition)?;
                self.process_statement(&mut if_stmt.then_branch)?;
                if let Some(else_branch) = &mut if_stmt.else_branch {
                    self.process_statement(else_branch)?;
                }
                Ok(())
            }
            Statement::For(for_stmt) => {
                match &mut for_stmt.variant {
                    ForVariant::Traditional {
                        initializer,
                        condition,
                        increment,
                    } => {
                        if let Some(init) = initializer {
                            self.process_statement(init)?;
                        }
                        if let Some(cond) = condition {
                            self.process_expression(cond)?;
                        }
                        if let Some(inc) = increment {
                            self.process_statement(inc)?;
                        }
                    }
                    ForVariant::Condition { condition } => {
                        self.process_expression(condition)?;
                    }
                    ForVariant::Range {
                        variable_type,
                        iterable,
                        ..
                    } => {
                        self.process_type(variable_type)?;
                        self.process_expression(iterable)?;
                    }
                }
                self.process_statement(for_stmt.body.as_mut())
            }
            Statement::While(while_stmt) => {
                self.process_expression(&mut while_stmt.condition)?;
                self.process_statement(&mut while_stmt.body)
            }
            Statement::Match(match_stmt) => {
                self.process_expression(&mut match_stmt.expr)?;
                for arm in &mut match_stmt.arms {
                    self.process_pattern(&mut arm.pattern)?;
                    if let Some(guard) = &mut arm.guard {
                        self.process_expression(guard)?;
                    }
                    self.process_statement(&mut arm.body)?;
                }
                Ok(())
            }
            Statement::Select(select_stmt) => {
                for case in &mut select_stmt.cases {
                    self.process_expression(&mut case.channel)?;
                    self.process_statement(&mut case.body)?;
                }
                if let Some(default_case) = &mut select_stmt.default_case {
                    self.process_statement(default_case)?;
                }
                Ok(())
            }
            Statement::Go(go_stmt) => self.process_call_expr(&mut go_stmt.call),
            Statement::Defer(defer_stmt) => self.process_expression(&mut defer_stmt.expr),
            Statement::Recover(recover_stmt) => {
                for arm in &mut recover_stmt.arms {
                    self.process_pattern(&mut arm.pattern)?;
                    if let Some(guard) = &mut arm.guard {
                        self.process_expression(guard)?;
                    }
                    self.process_statement(&mut arm.body)?;
                }
                Ok(())
            }
            Statement::Block(block_stmt) => {
                for stmt in &mut block_stmt.statements {
                    self.process_statement(stmt)?;
                }
                Ok(())
            }
            Statement::Export(_)
            | Statement::Break(_)
            | Statement::Continue(_)
            | Statement::Empty => Ok(()),
        }
    }

    fn process_assignment_target(&mut self, target: &mut AssignmentTarget) -> Result<()> {
        match target {
            AssignmentTarget::Variable(_) => Ok(()),
            AssignmentTarget::FieldAccess(field) => self.process_expression(&mut field.object),
            AssignmentTarget::IndexAccess(index) => {
                self.process_expression(&mut index.object)?;
                self.process_expression(&mut index.index)
            }
            AssignmentTarget::Deref(expr) => self.process_expression(expr),
        }
    }

    fn process_call_expr(&mut self, call: &mut CallExpr) -> Result<()> {
        self.process_expression(call.callee.as_mut())?;
        for arg in &mut call.arguments {
            self.process_expression(arg)?;
        }
        self.handle_generic_call(call)
    }

    fn process_expression(&mut self, expr: &mut Expression) -> Result<()> {
        match expr {
            Expression::Literal(_) | Expression::Variable(_) => Ok(()),
            Expression::Call(call) => self.process_call_expr(call),
            Expression::MethodCall(method_call) => {
                self.process_expression(method_call.object.as_mut())?;
                for arg in &mut method_call.arguments {
                    self.process_expression(arg)?;
                }
                if !method_call.type_args.is_empty() {
                    return Err(CompilerError::semantic(
                        method_call.location.line,
                        method_call.location.column,
                        "Generic method calls are not supported yet",
                    ));
                }
                Ok(())
            }
            Expression::Binary(binary) => {
                self.process_expression(binary.left.as_mut())?;
                self.process_expression(binary.right.as_mut())
            }
            Expression::Unary(unary) => self.process_expression(unary.operand.as_mut()),
            Expression::Try(try_expr) => self.process_expression(try_expr.expression.as_mut()),
            Expression::ForceUnwrap(force_expr) => {
                self.process_expression(force_expr.expression.as_mut())
            }
            Expression::Conditional(cond) => {
                self.process_expression(cond.condition.as_mut())?;
                self.process_expression(cond.true_expr.as_mut())?;
                self.process_expression(cond.false_expr.as_mut())
            }
            Expression::Array(array) => {
                for elem in &mut array.elements {
                    self.process_expression(elem)?;
                }
                Ok(())
            }
            Expression::Slice(slice) => {
                for elem in &mut slice.elements {
                    self.process_expression(elem)?;
                }
                self.process_type(&mut slice.element_type)?;
                if let Some(capacity) = &mut slice.capacity {
                    self.process_expression(capacity)?;
                }
                if let Some(count) = &mut slice.count {
                    self.process_expression(count)?;
                }
                Ok(())
            }
            Expression::Map(map) => {
                for entry in &mut map.entries {
                    self.process_expression(&mut entry.key)?;
                    self.process_expression(&mut entry.value)?;
                }
                Ok(())
            }
            Expression::Tuple(tuple) => {
                for elem in &mut tuple.elements {
                    self.process_expression(elem)?;
                }
                Ok(())
            }
            Expression::Struct(struct_expr) => {
                self.process_type(&mut struct_expr.struct_type)?;
                for field in &mut struct_expr.fields {
                    self.process_expression(&mut field.value)?;
                }
                Ok(())
            }
            Expression::FieldAccess(field) => self.process_expression(field.object.as_mut()),
            Expression::IndexAccess(index) => {
                self.process_expression(index.object.as_mut())?;
                self.process_expression(index.index.as_mut())
            }
            Expression::SliceRange(slice_range) => {
                self.process_expression(slice_range.object.as_mut())?;
                if let Some(start) = &mut slice_range.start {
                    self.process_expression(start)?;
                }
                if let Some(end) = &mut slice_range.end {
                    self.process_expression(end)?;
                }
                Ok(())
            }
            Expression::Cast(cast_expr) => {
                self.process_expression(cast_expr.expr.as_mut())?;
                self.process_type(&mut cast_expr.target_type)
            }
            Expression::Assert(assert_expr) => {
                self.process_expression(assert_expr.expr.as_mut())?;
                self.process_type(&mut assert_expr.asserted_type)
            }
            Expression::Lambda(lambda) => {
                for param in &mut lambda.parameters {
                    self.process_type(&mut param.param_type)?;
                }
                if let Some(ret_type) = &mut lambda.return_type {
                    self.process_type(ret_type)?;
                }
                self.process_expression(lambda.body.as_mut())
            }
            Expression::Match(match_expr) => {
                self.process_expression(match_expr.expr.as_mut())?;
                for arm in &mut match_expr.arms {
                    self.process_pattern(&mut arm.pattern)?;
                    if let Some(guard) = &mut arm.guard {
                        self.process_expression(guard)?;
                    }
                    self.process_expression(&mut arm.body)?;
                }
                Ok(())
            }
            Expression::Block(block_expr) => {
                for stmt in &mut block_expr.statements {
                    self.process_statement(stmt)?;
                }
                if let Some(final_expr) = &mut block_expr.final_expr {
                    self.process_expression(final_expr)?;
                }
                Ok(())
            }
            Expression::If(if_expr) => {
                self.process_expression(if_expr.condition.as_mut())?;
                self.process_expression(if_expr.then_branch.as_mut())?;
                if let Some(else_branch) = &mut if_expr.else_branch {
                    self.process_expression(else_branch)?;
                }
                Ok(())
            }
            Expression::For(for_expr) => {
                self.process_expression(for_expr.iterable.as_mut())?;
                self.process_expression(for_expr.body.as_mut())
            }
            Expression::While(while_expr) => {
                self.process_expression(while_expr.condition.as_mut())?;
                self.process_expression(while_expr.body.as_mut())
            }
            Expression::Select(select_expr) => {
                for case in &mut select_expr.cases {
                    self.process_expression(&mut case.channel)?;
                    self.process_expression(&mut case.body)?;
                }
                if let Some(default_case) = &mut select_expr.default_case {
                    self.process_expression(default_case)?;
                }
                Ok(())
            }
            Expression::Go(go_expr) => self.process_call_expr(&mut go_expr.call),
            Expression::New(new_expr) => {
                if let Some(initializer) = &mut new_expr.initializer {
                    self.process_expression(initializer)?;
                }
                Ok(())
            }
            Expression::Rc(rc_expr) => {
                self.process_type(&mut rc_expr.inner_type)?;
                if let Some(init) = &mut rc_expr.initializer {
                    self.process_expression(init)?;
                }
                Ok(())
            }
            Expression::Path(_) => Ok(()),
            Expression::FunctionLiteral(func_lit) => {
                if let Some(lambda) = &mut func_lit.lambda {
                    self.process_expression(lambda.body.as_mut())?;
                }
                if let Some(function_type) = &mut func_lit.function_type {
                    self.process_type(function_type)?;
                }
                Ok(())
            }
            Expression::Panic(panic_expr) => self.process_expression(panic_expr.value.as_mut()),
        }
    }

    fn handle_generic_call(&mut self, call: &mut CallExpr) -> Result<()> {
        if call.type_args.is_empty() {
            return Ok(());
        }

        self.ensure_concrete_type_args(&call.type_args, call.location)?;

        let new_name = match call.callee.as_mut() {
            Expression::Variable(name) => {
                let instantiated = self
                    .monomorphizer
                    .instantiate_function(name, &call.type_args)?;
                instantiated.name
            }
            Expression::Path(path) if path.segments.len() == 1 => {
                let instantiated = self
                    .monomorphizer
                    .instantiate_function(&path.segments[0], &call.type_args)?;
                instantiated.name
            }
            _ => {
                return Err(CompilerError::semantic(
                    call.location.line,
                    call.location.column,
                    "Generic type arguments are only supported on direct function calls",
                ));
            }
        };

        *call.callee = Expression::Variable(new_name);
        call.type_args.clear();
        call.resolved = None;
        Ok(())
    }

    fn ensure_concrete_type_args(&self, args: &[Type], location: Location) -> Result<()> {
        if let Some(arg) = args.iter().find(|ty| ty.contains_generic()) {
            return Err(CompilerError::semantic(
                location.line,
                location.column,
                format!(
                    "Cannot monomorphize call with unresolved generic type argument '{}'",
                    arg.to_string()
                ),
            ));
        }
        Ok(())
    }

    fn process_pattern(&mut self, pattern: &mut Pattern) -> Result<()> {
        match pattern {
            Pattern::Literal(_) | Pattern::Variable(_) | Pattern::Wildcard => Ok(()),
            Pattern::Tuple(patterns) | Pattern::Or(patterns) => {
                for pat in patterns {
                    self.process_pattern(pat)?;
                }
                Ok(())
            }
            Pattern::Struct(struct_pattern) => {
                self.process_type(&mut struct_pattern.struct_type)?;
                for field in &mut struct_pattern.fields {
                    self.process_pattern(&mut field.pattern)?;
                }
                Ok(())
            }
            Pattern::Array(patterns) => {
                for pat in patterns {
                    self.process_pattern(pat)?;
                }
                Ok(())
            }
            Pattern::Type(ty) => self.process_type(ty),
            Pattern::Enum(enum_pattern) => {
                for pat in &mut enum_pattern.patterns {
                    self.process_pattern(pat)?;
                }
                Ok(())
            }
        }
    }

    fn process_type(&mut self, ty: &mut Type) -> Result<()> {
        match ty {
            Type::Struct(struct_type) => {
                for arg in &mut struct_type.type_args {
                    self.process_type(arg)?;
                }
                if !struct_type.type_args.is_empty() {
                    self.ensure_concrete_type_args(&struct_type.type_args, struct_type.location)?;
                    let instantiated = self
                        .monomorphizer
                        .instantiate_struct(&struct_type.name, &struct_type.type_args)?;
                    struct_type.name = instantiated.name;
                    struct_type.type_args.clear();
                }
                Ok(())
            }
            Type::Reference(reference_type) => {
                self.process_type(reference_type.referenced_type.as_mut())
            }
            Type::Array(array_type) => self.process_type(array_type.element_type.as_mut()),
            Type::Map(map_type) => {
                self.process_type(map_type.key_type.as_mut())?;
                self.process_type(map_type.value_type.as_mut())
            }
            Type::Tuple(tuple_type) => {
                for element in &mut tuple_type.element_types {
                    self.process_type(element)?;
                }
                Ok(())
            }
            Type::Enum(enum_type) => {
                for arg in &mut enum_type.type_args {
                    self.process_type(arg)?;
                }
                Ok(())
            }
            Type::Function(function_type) => {
                for param in &mut function_type.parameter_types {
                    self.process_type(param)?;
                }
                if let Some(ret_type) = &mut function_type.return_type {
                    self.process_type(ret_type)?;
                }
                Ok(())
            }
            Type::Vec(vec_type) => self.process_type(vec_type.element_type.as_mut()),
            Type::Channel(channel_type) => self.process_type(channel_type.element_type.as_mut()),
            Type::Optional(optional_type) => self.process_type(optional_type.inner_type.as_mut()),
            Type::Error(error_type) => self.process_type(error_type.inner_type.as_mut()),
            Type::Result(result_type) => {
                self.process_type(result_type.ok_type.as_mut())?;
                self.process_type(result_type.err_type.as_mut())
            }
            Type::Option(option_type) => self.process_type(option_type.some_type.as_mut()),
            Type::Union(union_type) => {
                for member in &mut union_type.member_types {
                    self.process_type(member)?;
                }
                Ok(())
            }
            Type::Alias(alias_type) => {
                for arg in &mut alias_type.type_args {
                    self.process_type(arg)?;
                }
                if !alias_type.type_args.is_empty()
                    && self.monomorphizer.has_struct_template(&alias_type.name)
                {
                    self.ensure_concrete_type_args(&alias_type.type_args, alias_type.location)?;
                    let instantiated = self
                        .monomorphizer
                        .instantiate_struct(&alias_type.name, &alias_type.type_args)?;
                    *ty = Type::Struct(StructType {
                        name: instantiated.name,
                        type_args: vec![],
                        location: alias_type.location,
                    });
                }
                Ok(())
            }
            Type::Rc(rc_type) => self.process_type(rc_type.inner_type.as_mut()),
            Type::Weak(weak_type) => self.process_type(weak_type.inner_type.as_mut()),
            Type::Basic(_) | Type::Generic(_) => Ok(()),
        }
    }
}
